{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "635d8ebb",
   "metadata": {},
   "source": [
    "# ConversationKGMemory\n",
    "\n",
    "- Author: [Secludor](https://github.com/Secludor)\n",
    "- Design: [Secludor](https://github.com/Secludor)\n",
    "- Peer Review : [ulysyszh](https://github.com/ulysyszh), [Jinu Cho](https://github.com/jinucho)\n",
    "- Proofread : [Juni Lee](https://www.linkedin.com/in/ee-juni)\n",
    "- This is a part of [LangChain Open Tutorial](https://github.com/LangChain-OpenTutorial/LangChain-OpenTutorial)\n",
    "\n",
    "[![Open in Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/LangChain-OpenTutorial/LangChain-OpenTutorial/blob/main/05-Memory/05-ConversationKGMemory.ipynb)[![Open in GitHub](https://img.shields.io/badge/Open%20in%20GitHub-181717?style=flat-square&logo=github&logoColor=white)](https://github.com/LangChain-OpenTutorial/LangChain-OpenTutorial/blob/main/05-Memory/05-ConversationKGMemory.ipynb)\n",
    "## Overview\n",
    "\n",
    "Unlike ```ConversationEntityMemory```, which manages information about entities in a key-value format for individual entities, ```ConversationKGMemory```(Conversation Knowledge Graph Memory) is a module that manages relationships between entities in a graph format.\n",
    "\n",
    "It extracts and structures **knowledge triplets** (subject-relationship-object) to identify and store complex relationships between entities, and allows exploration of entity connectivity through **graph structure** .\n",
    "\n",
    "This helps the model understand relationships between different entities and better respond to queries based on complex networks and historical context.\n",
    "\n",
    "### Table of Contents\n",
    "\n",
    "- [Overview](#overview)\n",
    "- [Environment Setup](#environment-setup)\n",
    "- [Conversation Knowledge Graph Memory](#conversation-knowledge-graph-memory)\n",
    "- [Applying KG Memory to Chain](#applying-kg-memory-to-chain)\n",
    "- [Applying KG Memory with LCEL](#applying-kg-memory-with-lcel)\n",
    "\n",
    "### References\n",
    "\n",
    "- [LangChain Python API Reference>langchain-community: 0.3.13>memory>ConversationKGMemory](https://python.langchain.com/api_reference/community/memory/langchain_community.memory.kg.ConversationKGMemory.html)\n",
    "- [LangChain Python API Reference>langchain-community: 0.2.16>NetworkxEntityGraph](https://python.langchain.com/v0.2/api_reference/community/graphs/langchain_community.graphs.networkx_graph.NetworkxEntityGraph.html#langchain_community.graphs.networkx_graph.NetworkxEntityGraph)\n",
    "----"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c6c7aba4",
   "metadata": {},
   "source": [
    "## Environment Setup\n",
    "\n",
    "Set up the environment. You may refer to [Environment Setup](https://wikidocs.net/257836) for more details.\n",
    "\n",
    "**[Note]**\n",
    "- ```langchain-opentutorial``` is a package that provides a set of easy-to-use environment setup, useful functions and utilities for tutorials. \n",
    "- You can checkout the [```langchain-opentutorial```](https://github.com/LangChain-OpenTutorial/langchain-opentutorial-pypi) for more details."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "21943adb",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%capture --no-stderr\n",
    "%pip install langchain-opentutorial"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "f25ec196",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Install required packages\n",
    "from langchain_opentutorial import package\n",
    "\n",
    "package.install(\n",
    "    [\n",
    "        \"langsmith\",\n",
    "        \"langchain\",\n",
    "        \"langchain_core\",\n",
    "        \"langchain_community\",\n",
    "        \"langchain_openai\",\n",
    "    ],\n",
    "    verbose=False,\n",
    "    upgrade=False,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "7f9065ea",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Environment variables have been set successfully.\n"
     ]
    }
   ],
   "source": [
    "# Set environment variables\n",
    "from langchain_opentutorial import set_env\n",
    "\n",
    "set_env(\n",
    "    {\n",
    "        \"OPENAI_API_KEY\": \"\",\n",
    "        \"LANGCHAIN_API_KEY\": \"\",\n",
    "        \"LANGCHAIN_TRACING_V2\": \"true\",\n",
    "        \"LANGCHAIN_ENDPOINT\": \"https://api.smith.langchain.com\",\n",
    "        \"LANGCHAIN_PROJECT\": \"05-ConversationKGMemory\",  # title 과 동일하게 설정해 주세요\n",
    "    }\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "690a9ae0",
   "metadata": {},
   "source": [
    "You can alternatively set API keys such as ```OPENAI_API_KEY``` in a ```.env``` file and load them.\n",
    "\n",
    "[Note] This is not necessary if you've already set the required API keys in previous steps."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "4f99b5b6",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Load API keys from .env file\n",
    "from dotenv import load_dotenv\n",
    "\n",
    "load_dotenv(override=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "616661ad",
   "metadata": {},
   "source": [
    "## Conversation Knowledge Graph Memory\n",
    "\n",
    "```ConversationKGMemory``` is a memory module that stores and manages information extracted from conversations in a graph structure. \n",
    "\n",
    "This example demonstrates the following key features:\n",
    "- Storing conversation context (```save_context```)\n",
    "- (Reference) Getting a list of entity names in the graph sorted by causal dependence. (```get_topological_sort```)\n",
    "- Extracting entities from current conversation (```get_current_entities```)\n",
    "- Extracting knowledge triplets (```get_knowledge_triplets```)\n",
    "- Retrieving stored memory (```load_memory_variables```)\n",
    "\n",
    "The following example shows the process of extracting entities and relationships from a conversation about a new designer, Shelly Kim, and storing them in a graph format."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "17efec71",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_openai import ChatOpenAI\n",
    "from langchain_community.memory.kg import ConversationKGMemory"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "e165b797",
   "metadata": {},
   "outputs": [],
   "source": [
    "llm = ChatOpenAI(model_name=\"gpt-4o\", temperature=0)\n",
    "\n",
    "memory = ConversationKGMemory(llm=llm, return_messages=True)\n",
    "memory.save_context(\n",
    "    {\"input\": \"This is Shelly Kim who lives in Pangyo.\"},\n",
    "    {\"output\": \"Hello Shelly, nice to meet you! What kind of work do you do?\"},\n",
    ")\n",
    "memory.save_context(\n",
    "    {\"input\": \"Shelly Kim is our company's new designer.\"},\n",
    "    {\n",
    "        \"output\": \"That's great! Welcome to our team. I hope you'll enjoy working with us.\"\n",
    "    },\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dfed7692",
   "metadata": {},
   "source": [
    "### (Reference) get_topological_sort() → List[str]\n",
    "\n",
    "You can use the ```get_topological_sort``` method to view all entities stored in the knowledge graph in topological order:\n",
    "\n",
    "This method:\n",
    "- Uses ```NetworkX``` library to analyze the knowledge graph structure.\n",
    "- Performs topological sorting based on directed edges.\n",
    "- Returns a list of entities in dependency order.\n",
    "\n",
    "The order reflects the relationships between entities in the conversation, showing how they are connected in the knowledge graph."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "dc39b1f8",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['Shelly Kim', 'Pangyo', \"our company's new designer\"]"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "memory.kg.get_topological_sort()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4b2323cd",
   "metadata": {},
   "source": [
    "### get_current_entities(input_string: str) → List[str]\n",
    "\n",
    "Here's how the ```get_current_entities``` method works:\n",
    "\n",
    "**1. Entity Extraction Chain Creation**\n",
    "- Creates an ```LLMChain``` using the ```entity_extraction_prompt``` template.\n",
    "- This prompt is designed to extract proper nouns from the last line of the conversation.\n",
    "\n",
    "**2. Context Processing**\n",
    "- Retrieves the last **k*2** messages from the buffer. (default : k=2)\n",
    "- Generates conversation history string using ```human_prefix``` and ```ai_prefix```.\n",
    "\n",
    "**3. Entity Extraction**\n",
    "- Extracts proper nouns from the input string \"Who is Shelly Kim?\"\n",
    "- Primarily recognizes words starting with capital letters as proper nouns.\n",
    "- In this case, \"Shelly Kim\" is extracted as an entity.\n",
    "\n",
    "This method **only extracts entities from the question itself** , while the previous conversation context is used only for reference."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "1eae79db",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['Shelly Kim']"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "memory.get_current_entities({\"input\": \"Who is Shelly Kim?\"})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "28fe430b",
   "metadata": {},
   "source": [
    "### get_knowledge_triplets(input_string: str) → List[KnowledgeTriple]\n",
    "\n",
    "The ```get_knowledge_triplets``` method operates as follows:\n",
    "\n",
    "**1. Knowledge Triple Extraction Chain**\n",
    "- Creates an ```LLMChain``` using the ```knowledge_triplet_extraction_prompt``` template.\n",
    "- Designed to extract triples in ( **subject-relation-object** ) format from given text.\n",
    "\n",
    "**2. Memory Search**\n",
    "- Searches for information related to \"Shelly\" from previously stored conversations.\n",
    "- Stored context:\n",
    "  - \"This is Shelly Kim who lives in Pangyo.\"\n",
    "  - \"Shelly Kim is our company's new designer.\"\n",
    "\n",
    "**3. Triple Extraction**\n",
    "- Generates the following triples from the retrieved information:\n",
    "  - (Shelly Kim, lives in, Pangyo)\n",
    "  - (Shelly Kim, is, designer)\n",
    "  - (Shelly Kim, works at, our company)\n",
    "\n",
    "This method extracts relationship information in **triple format** from all stored conversation content **related to a specific entity** ."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "cb102317",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "([KnowledgeTriple(subject='Shelly Kim', predicate='lives in', object_='Pangyo'),\n",
       "  KnowledgeTriple(subject='Shelly Kim', predicate='is', object_=\"company's new designer\")],\n",
       " '\\n',\n",
       " [KnowledgeTriple(subject='Shelly Kim', predicate='lives in', object_='Pangyo')],\n",
       " '\\n',\n",
       " [KnowledgeTriple(subject='Shelly Kim', predicate='is a', object_='designer')],\n",
       " '\\n',\n",
       " [])"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "memory.get_knowledge_triplets({\"input\": \"Shelly\"}), \"\\n\", memory.get_knowledge_triplets(\n",
    "    {\"input\": \"Pangyo\"}\n",
    "), \"\\n\", memory.get_knowledge_triplets(\n",
    "    {\"input\": \"designer\"}\n",
    "), \"\\n\", memory.get_knowledge_triplets(\n",
    "    {\"input\": \"Langchain\"}\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "662b3831",
   "metadata": {},
   "source": [
    "### load_memory_variables(inputs: Dict[str, Any]) → Dict[str, Any]\n",
    "\n",
    "The ```load_memory_variables``` method operates through the following steps:\n",
    "\n",
    "**1. Entity Extraction**\n",
    "- Extracts entities (e.g., \"Shelly Kim\") from the input \"Who is Shelly Kim?\"\n",
    "- Internally uses the ```get_current_entities``` method.\n",
    "\n",
    "**2. Knowledge Retrieval**\n",
    "- Searches for all knowledge triplets related to the extracted entities.\n",
    "- Queries the graph for information previously stored via ```save_context``` method.\n",
    "\n",
    "**3. Information Formatting**\n",
    "- Converts found triplets into system messages.\n",
    "- Returns a list of message objects due to the ```return_messages=True``` setting.\n",
    "\n",
    "This method retrieves relevant information from the stored knowledge graph and returns it in a structured format, which can then be used as context for subsequent conversations with the language model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "92d6929d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'history': [SystemMessage(content=\"On Shelly Kim: Shelly Kim lives in Pangyo. Shelly Kim is our company's new designer.\", additional_kwargs={}, response_metadata={})]}"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "memory.load_memory_variables({\"input\": \"Who is Shelly Kim?\"})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aa00c3f4",
   "metadata": {},
   "source": [
    "## Applying KG Memory to Chain\n",
    "\n",
    "This section demonstrates how to use ```ConversationKGMemory``` with **ConversationChain** .\n",
    "\n",
    "(The class **ConversationChain** was deprecated in LangChain 0.2.7 and will be removed in 1.0. If you want, you can skip to [Applying KG Memory with LCEL](#applying-kg-memory-with-lcel))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "b020b140",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "C:\\Users\\Caelu\\AppData\\Local\\Temp\\ipykernel_5648\\1729312250.py:21: LangChainDeprecationWarning: The class `ConversationChain` was deprecated in LangChain 0.2.7 and will be removed in 1.0. Use :meth:`~RunnableWithMessageHistory: https://python.langchain.com/v0.2/api_reference/core/runnables/langchain_core.runnables.history.RunnableWithMessageHistory.html` instead.\n",
      "  conversation_with_kg = ConversationChain(\n"
     ]
    }
   ],
   "source": [
    "from langchain_community.memory.kg import ConversationKGMemory\n",
    "from langchain_core.prompts.prompt import PromptTemplate\n",
    "from langchain.chains import ConversationChain\n",
    "\n",
    "llm = ChatOpenAI(model_name=\"gpt-4o\", temperature=0)\n",
    "\n",
    "template = \"\"\"The following is a friendly conversation between a human and an AI. \n",
    "The AI is talkative and provides lots of specific details from its context. \n",
    "If the AI does not know the answer to a question, it truthfully says it does not know. \n",
    "The AI ONLY uses information contained in the \"Relevant Information\" section and does not hallucinate.\n",
    "\n",
    "Relevant Information:\n",
    "\n",
    "{history}\n",
    "\n",
    "Conversation:\n",
    "Human: {input}\n",
    "AI:\"\"\"\n",
    "prompt = PromptTemplate(input_variables=[\"history\", \"input\"], template=template)\n",
    "\n",
    "conversation_with_kg = ConversationChain(\n",
    "    llm=llm, prompt=prompt, memory=ConversationKGMemory(llm=llm)\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dcf33db4",
   "metadata": {},
   "source": [
    "Let's initialize the conversation with some basic information."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "e25d23b6",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "\"Hi Teddy! It's great to meet you. It sounds like you and Shelly are working together in a creative environment. Being a new designer, Shelly must be bringing fresh ideas and perspectives to your team. How has it been working with her so far?\""
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "conversation_with_kg.predict(\n",
    "    input=\"My name is Teddy. Shelly is a coworker of mine, and she's a new designer at our company.\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "de4d10e7",
   "metadata": {},
   "source": [
    "Let's query the memory for information about Shelly."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "c0e71b6e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'history': 'On Shelly: Shelly is a coworker of Teddy. Shelly is a new designer. Shelly works at our company.'}"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "conversation_with_kg.memory.load_memory_variables({\"input\": \"who is Shelly?\"})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bc034d9b",
   "metadata": {},
   "source": [
    "You can also reset the memory by ```memory.clear()```."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "208b5e4b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'history': ''}"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "conversation_with_kg.memory.clear()\n",
    "conversation_with_kg.memory.load_memory_variables({\"input\": \"who is Shelly?\"})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ba8d42f5",
   "metadata": {},
   "source": [
    "## Applying KG Memory with LCEL\n",
    "\n",
    "Let's examine the memory after having a conversation using a custom **ConversationChain** with ```ConversationKGMemory``` by LCEL"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "1905d380",
   "metadata": {},
   "outputs": [],
   "source": [
    "from operator import itemgetter\n",
    "from langchain_community.memory.kg import ConversationKGMemory\n",
    "from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder\n",
    "from langchain_core.runnables import RunnableLambda, RunnablePassthrough\n",
    "from langchain_openai import ChatOpenAI\n",
    "\n",
    "llm = ChatOpenAI(model_name=\"gpt-4o\", temperature=0)\n",
    "\n",
    "prompt = ChatPromptTemplate.from_messages(\n",
    "    [\n",
    "        (\n",
    "            \"system\",\n",
    "            \"\"\"The following is a friendly conversation between a human and an AI. \n",
    "The AI is talkative and provides lots of specific details from its context. \n",
    "If the AI does not know the answer to a question, it truthfully says it does not know. \n",
    "The AI ONLY uses information contained in the \"Relevant Information\" section and does not hallucinate.\n",
    "\n",
    "Relevant Information:\n",
    "{history}\"\"\",\n",
    "        ),\n",
    "        MessagesPlaceholder(variable_name=\"history\"),\n",
    "        (\"human\", \"{input}\"),\n",
    "    ]\n",
    ")\n",
    "\n",
    "memory = ConversationKGMemory(llm=llm, return_messages=True, memory_key=\"history\")\n",
    "\n",
    "\n",
    "class ConversationChain:\n",
    "    def __init__(self, prompt, llm, memory):\n",
    "        self.memory = memory\n",
    "        self.chain = (\n",
    "            RunnablePassthrough()\n",
    "            | RunnablePassthrough.assign(\n",
    "                history=RunnableLambda(memory.load_memory_variables)\n",
    "                | itemgetter(\"history\")\n",
    "            )\n",
    "            | prompt\n",
    "            | llm\n",
    "        )\n",
    "\n",
    "    def invoke(self, input_dict):\n",
    "        response = self.chain.invoke(input_dict)\n",
    "        self.memory.save_context(input_dict, {\"output\": response.content})\n",
    "        return response\n",
    "\n",
    "\n",
    "conversation_with_kg = ConversationChain(prompt, llm, memory)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "07ca6c16",
   "metadata": {},
   "source": [
    "Let's initialize the conversation with some basic information."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "99ac927b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "\"Hi Teddy! It's nice to meet you. It sounds like you and Shelly are working together at your company. How's everything going with the new designer on board?\""
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "response = conversation_with_kg.invoke(\n",
    "    {\n",
    "        \"input\": \"My name is Teddy. Shelly is a coworker of mine, and she's a new designer at our company.\"\n",
    "    }\n",
    ")\n",
    "response.content"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "514ea4d8",
   "metadata": {},
   "source": [
    "Let's query the memory for information about Shelly."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "c18820be",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'history': [SystemMessage(content='On Shelly: Shelly is a coworker of Teddy. Shelly is a new designer. Shelly works at our company.', additional_kwargs={}, response_metadata={})]}"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "conversation_with_kg.memory.load_memory_variables({\"input\": \"who is Shelly?\"})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a14d527d",
   "metadata": {},
   "source": [
    "You can also reset the memory by ```memory.clear()```."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "3dfac660",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'history': []}"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "conversation_with_kg.memory.clear()\n",
    "conversation_with_kg.memory.load_memory_variables({\"input\": \"who is Shelly?\"})"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "langchain-opentutorial-A2cWC-y3-py3.11",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
